{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert_multitask_learning import train_bert_multitask, train_eval_input_fn, BertMultiTask, params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/data3/yjp/bert-multitask-learning\n"
     ]
    }
   ],
   "source": [
    "cd ../"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Models\n",
    "If you don't want to control every thing, you can just call `train_bert_multitask` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saving preprocessing files to tmp\n",
      "INFO:tensorflow:Saving preprocessing files to tmp\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/weibo_cws_weibo_ner_ckpt', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a8e60e630>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a8e60e630>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f9a8e60e978>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "INFO:tensorflow:Create RestoreCheckpointHook.\n",
      "INFO:tensorflow:Skipping training since max_steps has already saved.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow_estimator.python.estimator.estimator.Estimator at 0x7f9a8e60ebe0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_bert_multitask(problem='weibo_ner&weibo_cws', num_gpus=1, num_epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to take more control of the training process, you can use lower level api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.estimator import Estimator\n",
    "from bert_multitask_learning.ckpt_restore_hook import RestoreCheckpointHook\n",
    "\n",
    "problem = 'weibo_ner&weibo_cws'\n",
    "num_gpus = 1\n",
    "bert_multitask_params = params.DynamicBatchSizeParams()\n",
    "\n",
    "# assign problem to params\n",
    "bert_multitask_params.train_epoch = 5\n",
    "bert_multitask_params.assign_problem(problem, gpu=num_gpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/weibo_cws_weibo_ner_ckpt', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a8e6064e0>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a8e6064e0>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f9a8e606c50>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "INFO:tensorflow:Create RestoreCheckpointHook.\n"
     ]
    }
   ],
   "source": [
    "# get model fn and create mirror strategy for distributed training\n",
    "model = BertMultiTask(params=bert_multitask_params)\n",
    "model_fn = model.get_model_fn()\n",
    "\n",
    "dist_trategy = tf.contrib.distribute.MirroredStrategy(\n",
    "    num_gpus=int(num_gpus),\n",
    "    cross_tower_ops=tf.contrib.distribute.AllReduceCrossDeviceOps(\n",
    "        'nccl', num_packs=int(num_gpus)))\n",
    "\n",
    "run_config = tf.estimator.RunConfig(\n",
    "    train_distribute=dist_trategy,\n",
    "    eval_distribute=dist_trategy,\n",
    "    log_step_count_steps=bert_multitask_params.log_every_n_steps)\n",
    "\n",
    "# create estimator\n",
    "estimator = Estimator(\n",
    "    model_fn,\n",
    "    model_dir=bert_multitask_params.ckpt_dir,\n",
    "    params=bert_multitask_params,\n",
    "    config=run_config)\n",
    "\n",
    "# pretrained bert restore hook\n",
    "train_hook = RestoreCheckpointHook(bert_multitask_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Skipping training since max_steps has already saved.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow_estimator.python.estimator.estimator.Estimator at 0x7f9a8e6065c0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# train\n",
    "def train_input_fn(): return train_eval_input_fn(bert_multitask_params)\n",
    "estimator.train(\n",
    "    train_input_fn, max_steps=bert_multitask_params.train_steps, hooks=[train_hook])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate and Predict\n",
    "\n",
    "For NER and CWS, we need different evaluation logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert_multitask_learning import eval_bert_multitask, predict_bert_multitask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/weibo_cws_weibo_ner_ckpt/', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a8f3e77b8>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a8f3e77b8>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f9a8f3e7160>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py:429: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "tf.py_func is deprecated in TF V2. Instead, use\n",
      "    tf.py_function, which takes a python function which manipulates tf eager\n",
      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
      "    being differentiable using a gradient tape.\n",
      "    \n",
      "INFO:tensorflow:Calling model_fn.\n",
      "WARNING:tensorflow:From /data3/yjp/bert-multitask-learning/bert_multitask_learning/bert/modeling.py:673: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dense instead.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/training/saver.py:1266: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use standard file APIs to check for files with this prefix.\n",
      "INFO:tensorflow:Restoring parameters from models/weibo_cws_weibo_ner_ckpt/model.ckpt-421\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing Inputs: 100%|██████████| 270/270 [00:00<00:00, 280.79it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'weibo_cws_Accuracy': 0.9494476297021396,\n",
       " 'weibo_cws_Accuracy Per Sequence': 0.35555555555555557}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_bert_multitask(problem='weibo_cws', model_dir='models/weibo_cws_weibo_ner_ckpt/', eval_scheme='acc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/weibo_cws_weibo_ner_ckpt/', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a440de898>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a440de898>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f9a440de588>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from models/weibo_cws_weibo_ner_ckpt/model.ckpt-421\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing Inputs: 100%|██████████| 270/270 [00:00<00:00, 281.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Acc Score: 0.981751\n",
      "Precision Score: 0.678899\n",
      "Recall Score: 0.691589\n",
      "F1 Score: 0.685185\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'Acc': 0.981750804083345,\n",
       " 'Precision': 0.6788990825688074,\n",
       " 'Recall': 0.6915887850467289,\n",
       " 'F1': 0.6851851851851852}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_bert_multitask(problem='weibo_ner', model_dir='models/weibo_cws_weibo_ner_ckpt/', eval_scheme='ner')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Checkpoint dir: models/weibo_cws_weibo_ner_ckpt\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_CPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:0\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:XLA_GPU:3\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:1\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:2\n",
      "INFO:tensorflow:Device is available but not used by distribute strategy: /device:GPU:3\n",
      "INFO:tensorflow:Configured nccl all-reduce.\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'models/weibo_cws_weibo_ner_ckpt', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a746dcd30>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.contrib.distribute.python.mirrored_strategy.MirroredStrategy object at 0x7f9a746dcd30>, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f9a746dc978>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from models/weibo_cws_weibo_ner_ckpt/model.ckpt-421\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing Inputs: 100%|██████████| 1/1 [00:00<00:00, 564.21it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[PAD]' 'B-GPE' 'I-GPE' 'O' 'B-GPE' 'I-GPE' 'O' 'O' 'O' 'O' 'O' '[PAD]']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# predict\n",
    "import numpy as np\n",
    "from bert_multitask_learning.utils import get_or_make_label_encoder\n",
    "predict_params = params.DynamicBatchSizeParams()\n",
    "# get prediction generator\n",
    "pred_prob = predict_bert_multitask(inputs=['中国和美国在打贸易战'], problem='weibo_cws&weibo_ner', params=predict_params)\n",
    "# get label encoder\n",
    "ner_label_encoder = get_or_make_label_encoder(params=predict_params, problem='weibo_ner', mode='predict')\n",
    "cws_label_encoder = get_or_make_label_encoder(params=predict_params, problem='weibo_cws', mode='predict')\n",
    "\n",
    "for prob in pred_prob:\n",
    "    ner_pred = np.argmax(prob['weibo_ner'], axis = -1)\n",
    "    print(ner_label_encoder.inverse_transform(ner_pred.tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export Model for Serving\n",
    "\n",
    "You can export the trained model for [serving](https://github.com/JayYip/bert-as-service)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert_multitask_learning import export_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:build graph...\n",
      "INFO:tensorflow:load parameters from checkpoint...\n",
      "INFO:tensorflow:freeze...\n",
      "WARNING:tensorflow:From /data3/yjp/bert-multitask-learning/bert_multitask_learning/export_model.py:77: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.compat.v1.graph_util.convert_variables_to_constants\n",
      "WARNING:tensorflow:From /data3/yjp/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/graph_util_impl.py:245: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.compat.v1.graph_util.extract_sub_graph\n",
      "INFO:tensorflow:Froze 201 variables.\n",
      "INFO:tensorflow:Converted 201 variables to const ops.\n",
      "INFO:tensorflow:write graph to: models/weibo_cws_weibo_ner_ckpt/export_model\n"
     ]
    }
   ],
   "source": [
    "export_model(bert_multitask_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
